package com.example.demo.websocket;

import com.fasterxml.jackson.databind.ObjectMapper;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.TextWebSocketHandler;

import java.io.IOException;
import java.util.Map;
import java.util.HashMap;
import java.util.concurrent.ConcurrentHashMap;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.annotation.PostConstruct;

@Component
public class WebSocketHandler extends TextWebSocketHandler {

    private static final Logger logger = LoggerFactory.getLogger(WebSocketHandler.class);
    private final Map<String, WebSocketSession> sessions = new ConcurrentHashMap<>();
    private final ObjectMapper objectMapper = new ObjectMapper();

    @PostConstruct
    public void init() {
        logger.info("WebSocket 处理器已初始化");
    }

    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        String openid = extractOpenidFromSession(session);
        sessions.put(openid, session);
        logger.info("新的连接已打开，openid: {}", openid);
    }

    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) throws Exception {
        String openid = extractOpenidFromSession(session);
        sessions.remove(openid);
        logger.info("连接已关闭，openid: {}", openid);
    }

    @Override
    protected void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception {
        String openid = extractOpenidFromSession(session);
        String payload = message.getPayload();
        logger.info("收到来自 openid: {} 的客户端消息: {}", openid, payload);

        try {
            Map<String, Object> messageMap = objectMapper.readValue(payload, Map.class);
            String type = (String) messageMap.get("type");

            if ("heartbeat".equals(type)) {
                handleHeartbeat(session, openid);
            } else if ("auth".equals(type)) {
                handleAuth(session, messageMap);
            } else {
                // 处理其他类型的消息
                logger.info("收到未知类型的消息: {}", type);
            }
        } catch (IOException e) {
            logger.error("解析消息失败: " + e.getMessage());
        }
    }

    private void handleHeartbeat(WebSocketSession session, String openid) throws IOException {
        Map<String, Object> response = new HashMap<>();
        response.put("type", "heartbeat");
        response.put("status", "ok");
        session.sendMessage(new TextMessage(objectMapper.writeValueAsString(response)));
        logger.info("响应心跳给 openid: {}", openid);
    }

    private void handleAuth(WebSocketSession session, Map<String, Object> messageMap) throws IOException {
        String userId = (String) messageMap.get("userId");
        String openid = (String) messageMap.get("openid");
        // 这里可以添加认证逻辑
        logger.info("用户认证: userId={}, openid={}", userId, openid);
        Map<String, Object> response = new HashMap<>();
        response.put("type", "auth");
        response.put("status", "success");
        session.sendMessage(new TextMessage(objectMapper.writeValueAsString(response)));
    }

    public void sendMessageToUser(String openid, String message) {
        WebSocketSession session = sessions.get(openid);
        if (session != null && session.isOpen()) {
            try {
                session.sendMessage(new TextMessage(message));
                logger.info("消息已发送给 openid: {}", openid);
            } catch (IOException e) {
                logger.error("发送消息给用户时发生错误: " + e.getMessage());
            }
        } else {
            logger.warn("未找到 openid: {} 的活跃连接", openid);
        }
    }

    public void broadcast(String message) {
        TextMessage textMessage = new TextMessage(message);
        for (Map.Entry<String, WebSocketSession> entry : sessions.entrySet()) {
            try {
                WebSocketSession session = entry.getValue();
                if (session.isOpen()) {
                    session.sendMessage(textMessage);
                    logger.info("广播消息已发送给 openid: {}", entry.getKey());
                }
            } catch (IOException e) {
                logger.error("广播消息时发生错误: " + e.getMessage());
            }
        }
    }

    private String extractOpenidFromSession(WebSocketSession session) {
        String query = session.getUri().getQuery();
        if (query != null && query.contains("openid=")) {
            return query.substring(query.indexOf("openid=") + 7);
        }
        return "unknown";
    }
}